package com.smh.aiagent.advisers;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.ai.document.Document;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Map;

public class MyLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    private static final Logger logger = LoggerFactory.getLogger(MyLoggerAdvisor.class);

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return Integer.MAX_VALUE - 1;
    } // 降低优先级,保证看到拼接好的输出

    public void before(AdvisedRequest advisedRequest) {

        Map<String, Object> map = advisedRequest.adviseContext();
        String userText = advisedRequest.userText();


        // 提取模板变量并替换
        Object docsObj = map.get("qa_retrieved_documents");
        if (docsObj instanceof ArrayList<?> docsList && !docsList.isEmpty()) {
            StringBuilder stringBuilder = new StringBuilder();
            for (Object doc : docsList) {
                if (doc instanceof Document document) {
                    stringBuilder.append(document.getText()).append("\n");
                }
            }
            String context = stringBuilder.toString();
            if (!context.isEmpty()) {
                userText = userText.replace("{question_answer_context}", context);
            }
        }

        logger.info("AI Request: {}", userText);
    }

    public void after(AdvisedResponse advisedResponse) {
        logger.info("AI Response: {}", advisedResponse.response().getResult().getOutput().getText());
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

        before(advisedRequest);

        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);

        after(advisedResponse);

        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

        before(advisedRequest);

        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);

        return new MessageAggregator().aggregateAdvisedResponse(advisedResponses,
                this::after);
    }
}